#include "heap.h"

#include <stdlib.h>
#include <stdio.h>

#define DEBUG 0
#if (DEBUG)
#define PRINTF(...) printf(__VA_ARGS__)
#define PRINTF_ARR(arr, n) do {           \
    for(int i = 0;i < n;i++)              \
    {                                     \
        PRINTF("%d ", arr[i + 1].value);  \
    }PRINTF("\n");                        \
}while(0)
#else
#define PRINTF(...)
#define PRINTF_ARR(...)
#endif

struct element_t {
    int value;
};

struct heap_t {
    unsigned int capacity;     //堆的容量
    unsigned int used;         //当前堆的元素个数
    struct element_t *element; //使用数组存储堆元素
    int (*compare)(int, int);  //比较器
};

static inline void heap_swap(int *value1, int *value2)
{
    if(*value1 == *value2) { return ; }

    *value1 = *value1 + *value2;
    *value2 = *value1 - *value2;
    *value1 = *value1 - *value2;
}

static inline int heap_compare(int value1, int value2)
{
    return value2 - value1;
}

static inline void heap_shiftup(struct heap_t *heap, unsigned int index)
{
    //存在父节点，并且满足向上堆化条件
    if((index >> 1) && (heap->compare(heap->element[index].value, heap->element[index >> 1].value) < 0))
    {
        heap_swap(&heap->element[index].value, &heap->element[index >> 1].value);
        heap_shiftup(heap, index >> 1);
    }
}

static inline void heap_shiftdown(struct heap_t *heap, unsigned int index, unsigned int size)
{
    unsigned int pos;

    pos = index;
    //存在左子节点，并且满足向下堆化条件
    if(((index << 1) <= size) && (heap->compare(heap->element[index].value, heap->element[index << 1].value) > 0))
    {
        pos = index << 1;
    }
    //存在右子节点，并且满足向下堆化条件
    if((((index << 1) + 1) <= size) && (heap->compare(heap->element[pos].value, heap->element[(index << 1) + 1].value) > 0))
    {
        pos = (index << 1) + 1;
    }
    if(pos == index)
    {
        return ;
    }
    heap_swap(&heap->element[index].value, &heap->element[pos].value);
    heap_shiftdown(heap, pos, size);
}

int heap_init(struct heap_t **heap, unsigned int capacity, int (*compare)(int, int))
{
    struct heap_t *new_heap;

    new_heap = malloc(sizeof(struct heap_t));
    if(!new_heap)
    {
        return -1;
    }
    new_heap->element = malloc(sizeof(struct element_t) * (capacity + 1));
    if(!new_heap->element)
    {
        free(new_heap);
        return -1;
    }
    new_heap->capacity = capacity;
    new_heap->used = 0;
    new_heap->compare = heap_compare;
    if(compare)
    {
        new_heap->compare = compare;
    }

    *heap = new_heap;

    return 0;
}

unsigned int heap_get_usedsize(struct heap_t *heap)
{
    return heap->used;
}

inline int heap_is_full(struct heap_t *heap)
{
    return heap->used == heap->capacity;
}

inline int heap_is_empty(struct heap_t *heap)
{
    return heap->used == 0;
}

int heap_peek_top(struct heap_t *heap, int *value)
{
    if(heap_is_empty(heap)) { return -1; }

    *value = heap->element[1].value;
    return 0;
}

int heap_peek_button(struct heap_t *heap, int *value)
{
    if(heap_is_empty(heap)) { return -1; }

    *value = heap->element[heap->used].value;
    return 0;
}

int heap_pop(struct heap_t *heap, int *value)
{
    if(heap_is_empty(heap)) { return -1; }

    *value = heap->element[1].value;
    return 0;
}

int heap_push(struct heap_t *heap, int value)
{
    heap->element[1].value = value;
    PRINTF("push start:");
    PRINTF_ARR(heap->element, heap->used);
    heap_shiftdown(heap, 1, heap->used);
    PRINTF("push end:");
    PRINTF_ARR(heap->element, heap->used);

    return 0;
}

int heap_sort(struct heap_t *heap)
{
    if(heap_is_empty(heap)) { return -1; }
    unsigned int index;

    PRINTF("sort start:");
    PRINTF_ARR(heap->element, heap->used);
    index = heap->used;
    do {
        //交换堆顶元素和堆底元素
        heap_swap(&heap->element[1].value, &heap->element[index].value);
        --index;
        //向下堆化
        heap_shiftdown(heap, 1, index);
    }while(index);
    PRINTF("sort end:");
    PRINTF_ARR(heap->element, heap->used);

    return 0;
}

int heap_insert(struct heap_t *heap, int value)
{
    if(heap_is_full(heap)) { return -1; }

    ++heap->used;
    //插入元素到堆底
    heap->element[heap->used].value = value;
    PRINTF("heap start:");
    PRINTF_ARR(heap->element, heap->used);
    //向上堆化
    heap_shiftup(heap, heap->used);
    PRINTF("heap end:");
    PRINTF_ARR(heap->element, heap->used);

    return 0;
}

int heap_remove(struct heap_t *heap, int *value)
{
    if(heap_is_empty(heap)) { return -1; }

    *value = heap->element[1].value;
    //交换堆顶元素和堆底元素
    heap_swap(&heap->element[1].value, &heap->element[heap->used].value);
    PRINTF("remove start:");
    PRINTF_ARR(heap->element, heap->used);
    --heap->used;
    //向下堆化
    heap_shiftdown(heap, 1, heap->used);
    PRINTF("remove end:");
    PRINTF_ARR(heap->element, heap->used);

    return 0;
}

void heap_destroy(struct heap_t *heap)
{
    free(heap->element);
    free(heap);
}

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
int* getLeastNumbers(int* arr, int arrSize, int k, int* returnSize){
    struct heap_t *heap;
    int *ret_arr;
    int value;
    int res;

    heap_init(&heap, k, NULL);
    for(int i = 0;i < k;i++)
    {
        heap_insert(heap, arr[i]);
    }
    for(int i = k;i < arrSize;i++)
    {
        res = heap_peek_top(heap, &value);
        if(res >= 0)
        {
            PRINTF("%d,%d\n", value, arr[i]);
            if(value > arr[i])
            {
                heap_pop(heap, &value);
                heap_push(heap, arr[i]);
            }
        }
    }
    heap_sort(heap);
    ret_arr = malloc(k * sizeof(*ret_arr));
    for(int i = 0;i < k;i++)
    {
        ret_arr[i] = heap->element[i + 1].value;
    }
    heap_destroy(heap);

    *returnSize = k;
    return ret_arr;
}
